package com.cloudea.learn.layer;

import com.cloudea.learn.Tensor;

/**
 * 矩阵向量乘法
 * y_j = sum_i (w_ji * x_i)
 * @author cloudea
 * @date 2022/01/02
 */
public class MatDot extends Layer {


    @Override
    protected Tensor forward(Context context, Tensor... inputs) {
        Tensor x = inputs[0];
        Tensor v = inputs[1];
        return Operator.gemv(x, v);
    }

    @Override
    public Tensor[] backward(Context context, Tensor grid, Tensor... inputs) {
        Tensor x = inputs[0];
        Tensor v = inputs[1];
        int m = x.getShape()[0];
        int n = x.getShape()[1];
        Tensor gridX = new Tensor(m, n);
        for (int i = 0; i < n; i++){
            for (int j = 0; j < m; j++){
                gridX.set(grid.get(j) * v.get(i), j, i);
            }
        }

        Tensor gridV = new Tensor(n);
        for (int i = 0; i < n; i++){
            double g = 0;
            for (int j = 0; j < m; j++){
                g += grid.get(j) * x.get(j, i);
            }
            gridV.set(g,i);
        }

        return new Tensor[]{gridX, gridV};
    }
}
